{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HMM Prediction "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import logging\n",
    "import itertools\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from hmmlearn.hmm import GaussianHMM\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tqdm import tqdm\n",
    "from docopt import docopt\n",
    "import tushare as ts\n",
    "import os\n",
    "\n",
    "class StockPredictor(object):\n",
    "    def __init__(self, company, test_size=0.8,\n",
    "                 n_hidden_states=2, n_latency_days=2,\n",
    "                 n_steps_frac_change=10, n_steps_frac_high=10,\n",
    "                 n_steps_frac_low=10):\n",
    "        self._init_logger()\n",
    " \n",
    "        self.company = company\n",
    "        if type(company) == str:\n",
    "            compnay = company\n",
    "        if type(company) == int:\n",
    "            company = str(company)\n",
    "        \n",
    "        self.data = ts.get_hist_data(str(company)).reset_index()\n",
    "        \n",
    "        self.n_latency_days = n_latency_days\n",
    "        self.hmm = GaussianHMM(n_components=n_hidden_states)\n",
    "        self._split_train_test_data(company, test_size)\n",
    " \n",
    "        self._compute_all_possible_outcomes(\n",
    "            n_steps_frac_change, n_steps_frac_high, n_steps_frac_low)\n",
    " \n",
    "    def _init_logger(self):\n",
    "        self._logger = logging.getLogger(__name__)\n",
    "        handler = logging.StreamHandler()\n",
    "        formatter = logging.Formatter(\n",
    "            '%(asctime)s %(name)-12s %(levelname)-8s %(message)s')\n",
    "        handler.setFormatter(formatter)\n",
    "        self._logger.addHandler(handler)\n",
    "        self._logger.setLevel(logging.DEBUG)\n",
    " \n",
    "    def _split_train_test_data(self,company, test_size):\n",
    "        data = ts.get_hist_data(company).reset_index()\n",
    "        test_data, _train_data = train_test_split(\n",
    "            data, test_size=test_size, shuffle=False)\n",
    " \n",
    "        self._train_data = _train_data\n",
    "        self._test_data = test_data\n",
    " \n",
    "    @staticmethod\n",
    "    def _extract_features(data):\n",
    "        open_price = np.array(data['open'])\n",
    "        close_price = np.array(data['close'])\n",
    "        high_price = np.array(data['high'])\n",
    "        low_price = np.array(data['low'])\n",
    " \n",
    "        # 计算收盘价、高价和低价的分数变化\n",
    "        # 这会用到一个特征\n",
    "        frac_change = (close_price - open_price) / open_price\n",
    "        frac_high = (high_price - open_price) / open_price\n",
    "        frac_low = (open_price - low_price) / open_price\n",
    " \n",
    "        return np.column_stack((frac_change, frac_high, frac_low))\n",
    " \n",
    "    def fit(self):\n",
    "        self._logger.info('>>> Extracting Features')\n",
    "        feature_vector = StockPredictor._extract_features(self._train_data)\n",
    "        self._logger.info('Features extraction Completed <<<')\n",
    " \n",
    "        self.hmm.fit(feature_vector)\n",
    " \n",
    "    def _compute_all_possible_outcomes(self, n_steps_frac_change,\n",
    "                                       n_steps_frac_high, n_steps_frac_low):\n",
    "        frac_change_range = np.linspace(-0.1, 0.1, n_steps_frac_change)\n",
    "        frac_high_range = np.linspace(0, 0.1, n_steps_frac_high)\n",
    "        frac_low_range = np.linspace(0, 0.1, n_steps_frac_low)\n",
    " \n",
    "        self._possible_outcomes = np.array(list(itertools.product(\n",
    "            frac_change_range, frac_high_range, frac_low_range)))\n",
    " \n",
    "    def _get_most_probable_outcome(self, day_index):\n",
    "        previous_data_start_index = max(0, day_index - self.n_latency_days)\n",
    "        previous_data_end_index = max(0, day_index - 1)\n",
    "        previous_data = self._test_data.iloc[previous_data_end_index: previous_data_start_index]\n",
    "        previous_data_features = StockPredictor._extract_features(\n",
    "            previous_data)\n",
    " \n",
    "        outcome_score = []\n",
    "        for possible_outcome in self._possible_outcomes:\n",
    "            total_data = np.row_stack(\n",
    "                (previous_data_features, possible_outcome))\n",
    "            outcome_score.append(self.hmm.score(total_data))\n",
    "        most_probable_outcome = self._possible_outcomes[np.argmax(\n",
    "            outcome_score)]\n",
    " \n",
    "        return most_probable_outcome\n",
    " \n",
    "    def predict_close_price(self, day_index):\n",
    "        open_price = self._test_data.iloc[day_index]['open']\n",
    "        predicted_frac_change, _, _ = self._get_most_probable_outcome(\n",
    "            day_index)\n",
    "        return open_price * (1 + predicted_frac_change)\n",
    " \n",
    "    def predict_close_prices_for_days(self, days, with_plot=False):\n",
    "        predicted_close_prices = []\n",
    "        for day_index in tqdm(range(days)):\n",
    "            predicted_close_prices.append(self.predict_close_price(day_index))\n",
    " \n",
    "        if with_plot:\n",
    "            test_data = self._test_data[0: days]\n",
    "            days = np.array(test_data['date'], dtype=\"datetime64[ms]\")\n",
    "            actual_close_prices = test_data['close']\n",
    " \n",
    "            fig = plt.figure()\n",
    " \n",
    "            axes = fig.add_subplot(111)\n",
    "            axes.plot(days, actual_close_prices, 'bo-', label=\"actual\")\n",
    "            axes.plot(days, predicted_close_prices, 'r+-', label=\"predicted\")\n",
    "            axes.set_title('{company}'.format(company=self.company)+\"HMM Prediciton\")\n",
    " \n",
    "            fig.autofmt_xdate()\n",
    " \n",
    "            plt.legend()\n",
    "            '''\n",
    "            if os.getcwd == 'img/':\n",
    "                print(\"CWD is right\")\n",
    "            if os.getcwd != 'img/':\n",
    "                os.chdir('img/')\n",
    "            '''\n",
    "            plt.savefig('{company}'.format(company=self.company)+\".png\")\n",
    "            plt.show()\n",
    " \n",
    "        return predicted_close_prices\n",
    "\n",
    "#stock_predictor = StockPredictor(601988)\n",
    "#stock_predictor.fit()\n",
    "#stock_predictor.predict_close_prices_for_days(10, with_plot=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
